# 创建一个简化版的 Memory，不涉及 Token 计数
from langchain.memory import ConversationTokenBufferMemory


class SimpleConversationTokenBufferMemory(ConversationTokenBufferMemory):
    def save_context(self, inputs: dict, outputs: dict) -> None:
        """重写保存上下文的方法，跳过 Token 检查"""
        input_str, output_str = self._get_input_output(inputs, outputs)

        # 直接添加到聊天历史，不检查 Token
        self.chat_memory.add_user_message(input_str)
        self.chat_memory.add_ai_message(output_str)